#-*- coding=utf-8 -*-
import numpy as np
import matplotlib.pyplot as mp
from matplotlib.animation import ArtistAnimation
import partie1_methodes as p1



def SolveMaillon1(theta, w, g, l, N, times, meth):
	"Fonction resolvant le problème du pendule simple"
	h = times/N
	eq = [(0.,[theta, w]),lambda X,t:np.array([X[1],(-g * np.sin(X[0]))/l])]
	return p1.meth_n_step(eq,N, h,p1.step_euler)
	

def PlotMaillon1(sol, N, times):
	"Fonction traçant la vitesse angulaire et theta en fonction du temps"
	h = times/N
	tms = np.arange(0.,times+h,h)
	
	sol_theta = map(lambda t : t[0], sol)
	sol_w = map(lambda t : t[1], sol)
	
	mp.clf()
	mp.plot(tms, sol_theta)
	mp.plot(tms, sol_w)
	mp.plot(tms,np.zeros(len(tms)),'k')
	mp.title("Variation de l'angle et de la vitesse du pendule a 1 maillon")
	mp.legend(("Angle theta","Vitesse angulaire"))
	mp.xlabel("Temps t")
	mp.ylabel("Angle theta")
	mp.savefig("../img/VariationTheta+VitesseMaillon1")
	mp.show()


def PlotTrajMaillon1(sol, l):
	"Fonction traçant la trajectoire du pendule simple"
	sol_theta = map(lambda t : t[0], sol)
	
	x = l*np.sin(sol_theta)
	y = -l*np.cos(sol_theta)
	
	mp.clf()
	mp.plot(0,0,marker='o', color='r', ls='')
	mp.plot(x, y)
	mp.plot([0,1/2.],[0,np.sqrt(3)/2.],'k')
	mp.title("Trajectoire du pendule a 1 maillon")
	mp.legend(("Trajectoire extremitee du pendule"))
	mp.xlabel("Temps t")
	mp.ylabel("Angle theta")
	mp.savefig("../img/TrajectoireMaillon1")
	mp.show()


def Maillon1():
	"Fonction resolvant puis traçant le problème du pendule simple"
	g = 9.8
	l = 1.
	
	N = 1240
	times = 20.
	
	theta = np.pi/4
	w = 0.

	meth = p1.step_euler

	sol = SolveMaillon1(theta, w, g, l, N, times, meth)
	PlotMaillon1(sol, N, times)
	PlotTrajMaillon1(sol, l)


def FreqMaillon1(theta):
	"Fonction determinant la fréquence d'oscilation du pendule simple"
	g = 9.8
	l = 1.
	
	N = 1240
	times = 20.
	h = times/N

	eq = [(0.,[theta,0.]), lambda X,t: np.array([X[1],(-g * np.sin(X[0]))/l])]

	x = np.arange(0.,times+h,h)
	
	sol = p1.meth_n_step(eq, N, h, p1.step_RK4)
	
	sol_theta = map(lambda t : t[0], sol)
	sol_w = map(lambda t : t[1], sol)

	tps = 0.
	Ts = []

	for i in range(0, len(sol_w)-1):
		if ((sol_w[i] < 0) and (sol_w[i+1] >= 0)):
			Ts.append(tps)
		tps += h

	T = (Ts[len(Ts)-1]-Ts[0]) / len(Ts)		
	return 1/T	

def FrecTraceMaillon1():
	"Fonction traçant le frequence d'oscilation d'un pendule simple pour differents theta0"
	theta = np.arange(-2*np.pi/3,2*np.pi/3,0.2)
	y = []
	for i in range(len(theta)):
		y.append(FreqMaillon1(theta[i]))
	mp.plot(theta, y, linewidth=1.0)
	mp.plot([theta[0],theta[len(theta)-1]],[np.sqrt(9.8/1.0)/(2*np.pi),np.sqrt(9.8/1.0)/(2*np.pi)])
	mp.xlabel("Angle initial")
	mp.ylabel("Frequence du pendule")
	mp.title("Frequence du pendule a 1 maillon")
	mp.savefig("../img/FrequenceMaillon1")
	mp.show()




## DEUX MAILLONS ##

def SolveMaillon2(theta1,theta2,w1,w2,g,l1,l2,m1,m2,N,times, meth):
	"Fonction resolvant le problème du pendule à deux maillons"
	h = times/N
	eq = [(0.,[theta1, w1, theta2, w2]), lambda X,T: np.array([X[1],(-g*(2*m1+m2)*np.sin(X[0])-m2*g*np.sin(X[0]-2*X[2])-2*np.sin(X[0]-X[2])*m2*(l2*(X[3]**2)+(X[1]**2)*l1*np.cos(X[0]-X[2])))/(l1*(2*m1 + m2 - m2*np.cos(2*X[0]-2*X[2]))), X[3], (2*np.sin(X[0]-X[2])*((X[1]**2)*l1*(m1+m2) + g*(m1+m2)*np.cos(X[0]) + (X[3]**2)*l2*m2*np.cos(X[0]-X[2]))) / (l2*(2*m1 + m2 - m2*np.cos(2*X[0]-2*X[2])))  ])]
	return p1.meth_n_step(eq, N, h, p1.step_RK4)
	

def PlotMaillon2(sol, N, times):
	"Fonction tracant l'angle et la vitesse angulaire de theta1 et theta2 en fonction du temps"
	h = times/N
	
	tms = np.arange(0.,times+h,h)
	
	sol_theta1 = map(lambda t : t[0], sol)
	sol_w1 = map(lambda t : t[1], sol)
	sol_theta2 = map(lambda t : t[2], sol)
	sol_w2 = map(lambda t : t[3], sol)
	
	mp.clf()
	mp.plot(tms, sol_theta1)
	mp.plot(tms, sol_theta2)
	#mp.plot(tms, sol_w1)
	#mp.plot(tms, sol_w2)
	
	mp.legend(("Angle theta1","Angle theta2","Vitesse angulaire 1","Vitesse angulaire 2"))
	mp.title("Modelisation du pendule a 2 maillons")
	mp.xlabel("Temps t")
	mp.ylabel("Angle theta")
	mp.savefig("../img/Maillon2Angle+Vitesse")
	mp.show()


def PlotTrajMaillon2(sol, l1, l2):
	"Fonction traçant la trajectoire du pendule à deux maillons"
	sol_theta1 = map(lambda t : t[0], sol)
	sol_theta2 = map(lambda t : t[2], sol)
	
	x1 = l1*np.sin(sol_theta1)
	y1 = -l1*np.cos(sol_theta1)
	
	x2 = l2*np.sin(sol_theta2)+x1
	y2 = -l2*np.cos(sol_theta2)+y1
	
	mp.clf()
	mp.plot(0,0,marker='o', color='r', ls='')
	mp.plot(x1, y1)
	mp.plot(x2, y2)
	mp.legend(("Trajectoire extremitee du pendule"))
	mp.title("Trajectoire du pendule a 2 maillons")
	mp.xlabel("Temps t")
	mp.ylabel("Angle theta")
	mp.savefig("../img/TrajectoireMaillon2")
	mp.show()





def AnimatonTrajMaillon(sol, l1, l2, N, times):
	"Fonction animant un ensemble de points correpondants ici au pendule."
	#Attention à ne pas quiter la fenetre d'animation avant la fin sous peine d'avoir une erreur.
	sol_theta1 = map(lambda t : t[0], sol)
	sol_theta2 = map(lambda t : t[2], sol)
	
	x1 = l1*np.sin(sol_theta1)
	y1 = -l1*np.cos(sol_theta1)
	
	x2 = l2*np.sin(sol_theta2)+x1
	y2 = -l2*np.cos(sol_theta2)+y1

	fig = mp.figure()
	ax  = fig.add_subplot(111) 

	ax.set_xlim(-20, 20)
	ax.set_ylim(-20,10)

	images = []

	for i in range(0, len(sol_theta2)):
	    line = ax.plot([0, x1[i]],[0, y1[i]], x2[0:i],y2[0:i], [x1[i], x2[i]],[y1[i], y2[i]], '-', color='b')
	    images.append((line))

	line_anim = ArtistAnimation(fig, images, interval=20, repeat=False)
	#line_anim.save('doublemaillon.mp4')

	mp.show()



def Maillon2Chaos():
	"Fonction resolvant le problème du pendule double et l'affichant, état chaotique"
	g = 9.8
	l1 = 10.
	l2 = 10.
	m1 = 1.
	m2 = 1.
	
	N = 300
	times = 40.
	
	theta1 = np.pi/2.
	theta2 = np.pi/3.
	w1 = np.sqrt(g/l2)
	w2 = 0.
	
	sol = SolveMaillon2(theta1, theta2, w1, w2, g, l1, l2, m1, m2, N, times, p1.step_RK4)
	PlotMaillon2(sol, N, times)
	PlotTrajMaillon2(sol, l1, l2)
	print "Chaotique par phenomène de resonance"
	print PremierRetournement(sol, N, times)
	AnimatonTrajMaillon(sol,l1, l2, N,times)


def Maillon2():
	"Fonction resolvant le problème du pendule double et l'affichant, etat non chaotique"
	g = 9.8
	l1 = 10.
	l2 = 10.
	m1 = 1.
	m2 = 1.
	
	N = 300
	times = 40.
	
	theta1 = np.pi/2.
	theta2 = np.pi/3.
	w1 = 0.
	w2 = 0.
	
	sol = SolveMaillon2(theta1, theta2, w1, w2, g, l1, l2, m1, m2, N, times, p1.step_RK4)
	PlotMaillon2(sol, N, times)
	PlotTrajMaillon2(sol, l1, l2)
	print "Non chaotique"
	print PremierRetournement(sol, N, times)


def PremierRetournement(sol, N, times):
	"Fonction calculant le temps de premier retournement (en unité de temps)"
	i = 0
	tps = 0.
	h = times/N
	sol_theta1 = map(lambda t : t[0], sol)
	sol_theta2 = map(lambda t : t[2], sol)
	
	for i in range(0, len(sol_theta2)):
		if(abs(sol_theta2[i]) <= np.pi):
			tps = tps + h
		else:
			return tps
	return "Pas de retournement"


def Maillon2InitCdt():
	"Fonction testant pour plusieurs theta2 initiaux proche le comportement du pendule"
	g = 9.8
	l1 = 10.
	l2 = 10.
	m1 = 1.
	m2 = 1.
	
	N = 300
	times = 40.
	
	theta1 = np.pi/2.
	theta2 = [np.pi/3.7,np.pi/3.4,np.pi/3.,np.pi/2.7,np.pi/2.5,np.pi/2.3,np.pi/2.,np.pi/1.8,np.pi/1.6]
	w1 = 0.
	w2 = 0.
	
	for i in range(0,len(theta2)):
		sol = SolveMaillon2(theta1, theta2[i], w1, w2, g, l1, l2, m1, m2, N, times, p1.step_RK4)
		print "Pendule double mailon theta2=",theta2[i]
		print "Premier retournement :", PremierRetournement(sol, N, times)
		AnimatonTrajMaillon(sol,l1, l2, N,times)





	
if __name__ == "__main__":    
	Maillon1()
	FrecTraceMaillon1()
	Maillon2()
	Maillon2Chaos()
	Maillon2InitCdt()





